import pandas as pd
import numpy as np
import torch
import ast

data = pd.read_csv("./MNIST/embedding_output_1k.csv")
data["embedding_vec"] = data['embedding'].apply(ast.literal_eval)
def consine_distance(a, b):
    # 计算两个向量的余弦夹角，夹角越接近1越相似，-1代表正好相反
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))


query = data['embedding_vec'][0]
distances = []
for vec in data['embedding_vec'].values[1:]:
    distances.append(consine_distance(query, vec))
indices = torch.topk(torch.tensor(distances), k=3).indices.data
print(f"indices:{indices}")

print(f"query:{data['combined'][0]}")
for i, index in enumerate(indices):
    print(f"{i}:{data['combined'][index.item()]}")
